#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
数据处理器模块
负责数据清洗、分类和匹配等数据处理操作
"""

import pandas as pd
import numpy as np
import logging
from typing import Optional, Callable

from utils.logger import get_logger

# 获取日志记录器
logger = get_logger(__name__)

class DataHandler:
    """数据处理类"""
    
    def __init__(self, config: dict = None):
        """
        初始化数据处理器
        
        Args:
            config (dict, optional): 配置字典
        """
        self.config = config or {}
        self.batch_size = self.config.get('processing', {}).get('batch_size', 1000)
    
    def process(self, df: pd.DataFrame, progress_callback: Optional[Callable] = None) -> pd.DataFrame:
        """
        主数据处理流程
        
        Args:
            df (pd.DataFrame): 输入数据框
            progress_callback (Callable, optional): 进度回调函数
            
        Returns:
            pd.DataFrame: 处理后的数据框
        """
        try:
            logger.info("开始数据处理流程")
            
            # 1. 数据清洗
            if progress_callback:
                progress_callback(10)
            cleaned_df = self._clean_data(df)
            
            # 2. 数据验证
            if progress_callback:
                progress_callback(30)
            validated_df = self._validate_data(cleaned_df)
            
            # 3. 数据转换和分类
            if progress_callback:
                progress_callback(50)
            transformed_df = self._transform_and_classify(validated_df)
            
            # 4. 数据匹配
            if progress_callback:
                progress_callback(70)
            matched_df = self._match_data(transformed_df)
            
            # 5. 数据后处理和格式化
            if progress_callback:
                progress_callback(90)
            final_df = self._post_process(matched_df)
            
            if progress_callback:
                progress_callback(100)
            
            logger.info("数据处理流程完成")
            return final_df
            
        except Exception as e:
            logger.error(f"数据处理失败: {str(e)}")
            raise
    
    def _clean_data(self, df: pd.DataFrame) -> pd.DataFrame:
        """
        数据清洗
        
        Args:
            df (pd.DataFrame): 输入数据框
            
        Returns:
            pd.DataFrame: 清洗后的数据框
        """
        logger.info("执行数据清洗")
        
        # 创建数据副本以避免修改原始数据
        cleaned_df = df.copy()
        
        # 删除完全重复的行
        initial_rows = len(cleaned_df)
        cleaned_df = cleaned_df.drop_duplicates()
        duplicate_rows = initial_rows - len(cleaned_df)
        if duplicate_rows > 0:
            logger.info(f"删除了 {duplicate_rows} 个重复行")
        
        # 处理缺失值
        # 对于数值列，用0填充
        numeric_columns = cleaned_df.select_dtypes(include=[np.number]).columns
        cleaned_df[numeric_columns] = cleaned_df[numeric_columns].fillna(0)
        
        # 对于字符串列，用空字符串填充
        string_columns = cleaned_df.select_dtypes(include=['object']).columns
        cleaned_df[string_columns] = cleaned_df[string_columns].fillna('')
        
        # 去除字符串列的前后空格
        for col in string_columns:
            cleaned_df[col] = cleaned_df[col].astype(str).str.strip()
        
        # 将列名标准化（去除空格，转为小写）
        cleaned_df.columns = cleaned_df.columns.str.strip().str.lower().str.replace(' ', '_')
        
        return cleaned_df
    
    def _validate_data(self, df: pd.DataFrame) -> pd.DataFrame:
        """
        数据验证
        
        Args:
            df (pd.DataFrame): 输入数据框
            
        Returns:
            pd.DataFrame: 验证后的数据框
        """
        logger.info("执行数据验证")
        
        validated_df = df.copy()
        
        # 这里可以添加具体的业务验证逻辑
        # 例如：检查必填字段、验证数据范围、格式等
        
        # 示例：标记可能有问题的数据行
        if '数量' in validated_df.columns:
            # 标记数量为负数的行
            validated_df['_数量异常'] = validated_df['数量'] < 0
            invalid_count = validated_df['_数量异常'].sum()
            if invalid_count > 0:
                logger.warning(f"发现 {invalid_count} 行数量异常数据")
        
        return validated_df
    
    def _transform_and_classify(self, df: pd.DataFrame) -> pd.DataFrame:
        """
        数据转换和分类
        
        Args:
            df (pd.DataFrame): 输入数据框
            
        Returns:
            pd.DataFrame: 转换和分类后的数据框
        """
        logger.info("执行数据转换和分类")
        
        transformed_df = df.copy()
        
        # 这里可以添加具体的数据转换和分类逻辑
        # 例如：创建新的计算列、应用分类规则等
        
        # 示例：如果有'长度'、'宽度'、'高度'列，计算体积
        dimensions = ['长度', '宽度', '高度']
        if all(dim in transformed_df.columns for dim in dimensions):
            transformed_df['体积'] = transformed_df['长度'] * transformed_df['宽度'] * transformed_df['高度']
            logger.info("计算了体积列")
        
        # 示例：基于某些条件进行分类
        if '重量' in transformed_df.columns:
            # 根据重量分类
            bins = [0, 10, 50, 100, float('inf')]
            labels = ['轻量级', '中量级', '重量级', '超重量级']
            transformed_df['重量分类'] = pd.cut(transformed_df['重量'], bins=bins, labels=labels)
            logger.info("完成重量分类")
        
        return transformed_df
    
    def _match_data(self, df: pd.DataFrame) -> pd.DataFrame:
        """
        数据匹配
        
        Args:
            df (pd.DataFrame): 输入数据框
            
        Returns:
            pd.DataFrame: 匹配后的数据框
        """
        logger.info("执行数据匹配")
        
        matched_df = df.copy()
        
        # 这里可以实现具体的数据匹配逻辑
        # 例如：基于某些关键字段进行匹配、排序等
        
        # 示例：如果有'产品代码'和'包装代码'列，可以根据这些列进行匹配处理
        # 这里只是占位，实际实现需要根据具体业务需求
        
        return matched_df
    
    def _post_process(self, df: pd.DataFrame) -> pd.DataFrame:
        """
        数据后处理和格式化
        
        Args:
            df (pd.DataFrame): 输入数据框
            
        Returns:
            pd.DataFrame: 后处理后的数据框
        """
        logger.info("执行数据后处理和格式化")
        
        final_df = df.copy()
        
        # 移除中间处理列（以'_'开头的列）
        temp_columns = [col for col in final_df.columns if col.startswith('_')]
        if temp_columns:
            final_df = final_df.drop(columns=temp_columns)
            logger.info(f"移除了 {len(temp_columns)} 个临时处理列")
        
        # 调整列的显示顺序（可以根据业务需求自定义）
        # 这里只是示例，实际顺序需要根据业务需求确定
        important_columns = ['订单号', '产品代码', '客户名称']
        existing_important = [col for col in important_columns if col in final_df.columns]
        
        if existing_important:
            # 重要列在前，其他列在后
            other_columns = [col for col in final_df.columns if col not in existing_important]
            final_df = final_df[existing_important + other_columns]
        
        # 格式化数值列（设置小数位数）
        decimal_places = self.config.get('excel', {}).get('decimal_places', 2)
        numeric_columns = final_df.select_dtypes(include=[np.number]).columns
        final_df[numeric_columns] = final_df[numeric_columns].round(decimal_places)
        
        return final_df
    
    def batch_process(self, df: pd.DataFrame, process_func: Callable, 
                     progress_callback: Optional[Callable] = None) -> pd.DataFrame:
        """
        批量处理数据
        
        Args:
            df (pd.DataFrame): 输入数据框
            process_func (Callable): 应用于每个批次的处理函数
            progress_callback (Callable, optional): 进度回调函数
            
        Returns:
            pd.DataFrame: 处理后的数据框
        """
        logger.info(f"开始批量处理，批次大小: {self.batch_size}")
        
        # 计算批次数
        n_batches = (len(df) + self.batch_size - 1) // self.batch_size
        results = []
        
        # 分批处理
        for i in range(n_batches):
            start_idx = i * self.batch_size
            end_idx = min((i + 1) * self.batch_size, len(df))
            
            # 提取批次数据
            batch = df.iloc[start_idx:end_idx].copy()
            
            # 处理批次
            processed_batch = process_func(batch)
            results.append(processed_batch)
            
            # 更新进度
            if progress_callback:
                progress = int(((i + 1) / n_batches) * 100)
                progress_callback(progress)
            
            logger.info(f"批次 {i+1}/{n_batches} 处理完成")
        
        # 合并结果
        final_df = pd.concat(results, ignore_index=True)
        logger.info("批量处理完成")
        
        return final_df